"""
Run Burgers equation experiments from Section 3 of the paper.

This is a simplified test version to verify the approach works.
"""

import sys
import os
import numpy as np
import time

# Add parent directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from locELM.core.domain import Domain, CollocationPoints
from locELM.core.networks import MultiSubdomainNetwork
from locELM.core.continuity import ContinuityConditions
from locELM.solvers.nonlinear import NonlinearPDESolver


def exact_solution_burgers(x, t):
    """Exact solution for the Burgers equation test case."""
    x = np.atleast_1d(x)
    t = np.atleast_1d(t)

    # Spatial component
    spatial = (1 + x/10) * (2*np.cos(np.pi*x + 2*np.pi/5) +
                            1.5*np.cos(2*np.pi*x - 3*np.pi/5))

    # Temporal component
    temporal = (1 + t/10) * (2*np.cos(np.pi*t + 2*np.pi/5) +
                             1.5*np.cos(2*np.pi*t - 3*np.pi/5))

    return spatial * temporal


def source_term_burgers(x, t, nu=0.01):
    """Source term (simplified - assumes zero for this test)."""
    return np.zeros_like(x)


def compute_burgers_operator(network, points, nu=0.01):
    """
    Compute Burgers operator: du/dt + u * du/dx - nu * d^2u/dx^2
    """
    eps = 1e-7
    n_points = points.shape[0]
    operator_values = np.zeros((n_points, 1))

    for i in range(n_points):
        pt = points[i:i+1, :]  # Shape (1, 2)

        # Compute u at point
        u = network(pt)[0, 0]

        # Compute du/dx using finite differences
        pt_xp = pt.copy()
        pt_xp[0, 0] += eps
        pt_xm = pt.copy()
        pt_xm[0, 0] -= eps
        u_xp = network(pt_xp)[0, 0]
        u_xm = network(pt_xm)[0, 0]
        du_dx = (u_xp - u_xm) / (2 * eps)

        # Compute d^2u/dx^2
        d2u_dx2 = (u_xp - 2*u + u_xm) / (eps**2)

        # Compute du/dt
        pt_tp = pt.copy()
        pt_tp[0, 1] += eps
        pt_tm = pt.copy()
        pt_tm[0, 1] -= eps
        u_tp = network(pt_tp)[0, 0]
        u_tm = network(pt_tm)[0, 0]
        du_dt = (u_tp - u_tm) / (2 * eps)

        # Burgers operator
        operator_values[i, 0] = du_dt + u * du_dx - nu * d2u_dx2

    return operator_values


def run_simple_burgers_test():
    """
    Run a simple Burgers equation test with a single time block.
    """
    print("\n" + "="*70)
    print("SIMPLE BURGERS EQUATION TEST (Single Time Block)")
    print("="*70)

    # Parameters
    a, b = 0.0, 5.0  # Spatial domain
    t_start, t_end = 0.0, 0.25  # Time domain (short for testing)
    nu = 0.01  # Viscosity
    n_x_subdomains = 3  # Fewer subdomains for testing
    q_x, q_t = 10, 10  # Fewer collocation points for testing
    n_hidden = 50  # Fewer training parameters for testing
    rand_mag = 0.75
    seed = 22

    np.random.seed(seed)

    print(f"Domain: x in [{a}, {b}], t in [{t_start}, {t_end}]")
    print(f"Spatial subdomains: {n_x_subdomains}")
    print(f"Collocation points per subdomain: {q_x} x {q_t}")
    print(f"Hidden nodes per subdomain: {n_hidden}")
    print(f"Viscosity: {nu}")

    # Create 2D domain (x, t)
    domain = Domain(
        bounds=[(a, b), (t_start, t_end)],
        n_subdomains=[n_x_subdomains, 1]  # Nx x 1 in time
    )

    # Create neural networks
    networks = MultiSubdomainNetwork(
        n_subdomains=domain.n_total,
        input_dim=2,
        output_dim=1,
        hidden_layers=[n_hidden],  # List with one layer
        activation='tanh',
        r_max=rand_mag,
        seed=seed
    )

    # Add networks for each subdomain
    for idx in range(domain.n_total):
        region = domain.get_subdomain_region(idx)
        networks.add_subdomain_network(region)

    # Create solver (which handles collocation points and continuity)
    solver = NonlinearPDESolver(
        domain=domain,
        networks=networks,
        continuity_order=[1, 0],  # C1 in x (dim 0), C0 in t (dim 1)
        collocation_type='uniform',
        n_collocation_points=[q_x, q_t]
    )

    # Define boundary conditions
    def bc_left(pts):
        """Left spatial boundary: x = a"""
        t_vals = pts[:, 1]
        x_vals = np.full_like(t_vals, a)
        return exact_solution_burgers(x_vals, t_vals)

    def bc_right(pts):
        """Right spatial boundary: x = b"""
        t_vals = pts[:, 1]
        x_vals = np.full_like(t_vals, b)
        return exact_solution_burgers(x_vals, t_vals)

    def bc_initial(pts):
        """Initial condition: t = t_start"""
        x_vals = pts[:, 0]
        t_vals = np.full_like(x_vals, t_start)
        return exact_solution_burgers(x_vals, t_vals)

    boundary_conditions = {
        0: {'min': bc_left, 'max': bc_right},  # x boundaries
        1: {'min': bc_initial}  # t initial condition
    }

    # Build nonlinear residual function
    def burgers_operator(network, points, subdomain_idx):
        return compute_burgers_operator(network, points, nu)

    def source_func(points):
        x_vals = points[:, 0]
        t_vals = points[:, 1]
        return source_term_burgers(x_vals, t_vals, nu)

    residual_function = solver.build_nonlinear_residual(
        nonlinear_operator=burgers_operator,
        source_term=source_func,
        boundary_conditions=boundary_conditions,
        output_dim=1
    )

    # Solve using NLSQ-perturb method
    print("\nSolving...")
    start_time = time.time()

    try:
        weights = solver.solve_nonlinear(
            residual_function=residual_function,
            method='nlsq-perturb',
            nlsq_options={
                'delta': 0.5,
                'cost_threshold': 1e-3,
                'max_sub_iterations': 20
            }
        )

        solve_time = time.time() - start_time
        print(f"Solve time: {solve_time:.2f} seconds")

        # Compute errors on evaluation grid
        print("\nComputing errors...")
        n_eval_x = 50
        n_eval_t = 10
        x_eval = np.linspace(a, b, n_eval_x)
        t_eval = np.linspace(t_start, t_end, n_eval_t)

        errors = []
        for t_val in t_eval:
            for x_val in x_eval:
                point = np.array([x_val, t_val])

                # Find subdomain
                subdomain_idx = domain.find_subdomain(point)
                network = networks.get_network(subdomain_idx)

                # Compute solution
                u_pred = network(point.reshape(1, -1))[0, 0]
                u_exact = exact_solution_burgers(np.array([x_val]), np.array([t_val]))[0]

                errors.append(abs(u_pred - u_exact))

        errors = np.array(errors)
        max_error = np.max(errors)
        rms_error = np.sqrt(np.mean(errors**2))

        print(f"Max error: {max_error:.3e}")
        print(f"RMS error: {rms_error:.3e}")

        # Save results
        output_dir = os.path.join(os.path.dirname(__file__), '..', 'experiments')
        os.makedirs(output_dir, exist_ok=True)

        np.savez(
            os.path.join(output_dir, 'burgers_simple_test.npz'),
            max_error=max_error,
            rms_error=rms_error,
            solve_time=solve_time,
            n_x_subdomains=n_x_subdomains,
            q_x=q_x,
            q_t=q_t,
            n_hidden=n_hidden
        )

        print(f"\nResults saved to experiments/burgers_simple_test.npz")

        return {
            'max_error': max_error,
            'rms_error': rms_error,
            'solve_time': solve_time
        }

    except Exception as e:
        print(f"\nError during solving: {e}")
        import traceback
        traceback.print_exc()
        return None


if __name__ == '__main__':
    result = run_simple_burgers_test()
    if result:
        print("\n" + "="*70)
        print("TEST COMPLETED SUCCESSFULLY")
        print("="*70)
